import os

import os

from torch.utils.data import Dataset,DataLoader

from torchvision.transforms import ToTensor
import torchvision.transforms as transforms
from PIL import Image
# import PIL
import numpy as np


# to tensor
# Converts a PIL Image or numpy.ndarray (H x W x C) in the range
#     [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0].


class GameDataset(Dataset):
    def __init__(self, img_dir, num_img=-1,transform=None, target_transform=None):
        # PIL.PngImagePlugin.MAX_TEXT_MEMORY = 6710886400
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

        self.img_paths = []
        self.file_names = os.listdir(self.img_dir)
        for file in self.file_names:
            if os.path.splitext(file)[1].endswith(('jpg', 'png')):
                self.img_paths.append("%s/%s" % (self.img_dir, file))
        self.img_paths=self.img_paths[:num_img]
        # print("img_paths",self.img_paths)

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):

        img_src = Image.open(self.img_paths[idx]).convert('RGB')
        if self.transform:
            img_src = self.transform(img_src)
        return img_src

